package weiyao.Detector;

import weiyao.GUI.ScanResultCallback;
import us.codecraft.webmagic.Page;
import us.codecraft.webmagic.selector.Html;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

public class WeakPasswordDetector extends BaseDetector {

    // 字段名识别列表
    private static final List<String> USERNAME_FIELDS = Arrays.asList(
            "username", "user", "login", "account", "email", "uname"
    );
    private static final List<String> PASSWORD_FIELDS = Arrays.asList(
            "password", "pass", "pwd", "passwd"
    );

    // 常见弱口令组合
    private static final String[] COMMON_CREDENTIALS = {
            "admin:admin",
            "admin:123456",
            "admin:password",
            "root:root",
            "root:123456",
            "test:test",
            "user:user",
            "administrator:password",
            "admin:admin123",
            "admin:qwerty"
    };

    // 爆破资源
    private final List<String> dictionary = new ArrayList<>();
    private final Set<String> testedUrls = Collections.synchronizedSet(new HashSet<>());
    private final AtomicBoolean isRunning = new AtomicBoolean(true);

    // 线程池和限流
    private ExecutorService executorService;
    private final Semaphore requestSemaphore = new Semaphore(5); // 限制并发请求数
    private final AtomicInteger activeRequests = new AtomicInteger(0);

    public WeakPasswordDetector() {
        // 加载爆破字典
        loadDictionary("C:\\Users\\12290\\Desktop\\爆破字典.txt");

        // 创建线程池
        int corePoolSize = Runtime.getRuntime().availableProcessors() * 2;
        executorService = Executors.newFixedThreadPool(corePoolSize, r -> {
            Thread t = new Thread(r);
            t.setDaemon(true); // 设置为守护线程，避免阻止JVM退出
            return t;
        });
    }

    public void stop() {
        isRunning.set(false);
        executorService.shutdownNow();
        try {
            if (!executorService.awaitTermination(3, TimeUnit.SECONDS)) {
                executorService.shutdownNow();
            }
        } catch (InterruptedException e) {
            executorService.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }

    private void loadDictionary(String filePath) {
        try (BufferedReader reader = new BufferedReader(new FileReader(filePath))) {
            String line;
            while ((line = reader.readLine()) != null) {
                line = line.trim();
                if (!line.isEmpty() && !dictionary.contains(line)) {
                    dictionary.add(line);
                }
            }
            System.out.println("字典加载完成，共 " + dictionary.size() + " 个条目");
        } catch (IOException e) {
            System.err.println("无法加载字典文件: " + e.getMessage());
        }
    }

    @Override
    public void detect(Page page, String baseUrl, ScanResultCallback callback) {


        Html html = page.getHtml();
        List<String> formElements = html.xpath("//form").all();

        for (String formHtml : formElements) {
            if (!isRunning.get()) return; // 如果停止扫描则退出

            Html form = new Html(formHtml);
            String action = form.xpath("//form/@action").get() != null ?
                    form.xpath("//form/@action").get() : baseUrl;
            String method = form.xpath("//form/@method").get() != null ?
                    form.xpath("//form/@method").get() : "POST";

            String formUrl = makeAbsoluteUrl(baseUrl, action);

            // 避免重复测试同一个URL
            if (testedUrls.contains(formUrl)) {
                continue;
            }
            testedUrls.add(formUrl);

            // 识别可能的用户名和密码字段
            Map<String, String> usernameFields = new HashMap<>();
            Map<String, String> passwordFields = new HashMap<>();

            List<String> inputElements = form.xpath("//input").all();
            for (String inputHtml : inputElements) {
                Html input = new Html(inputHtml);
                String name = input.xpath("//input/@name").get();
                String type = input.xpath("//input/@type").get();
                String value = input.xpath("//input/@value").get();

                if (name == null || name.isEmpty()) continue;

                // 识别用户名字段
                if (isUsernameField(name, type)) {
                    usernameFields.put(name, value != null ? value : "");
                }
                // 识别密码字段
                else if (isPasswordField(name, type)) {
                    passwordFields.put(name, value != null ? value : "");
                }
            }

            // 如果没有找到用户名或密码字段，跳过
            if (usernameFields.isEmpty() || passwordFields.isEmpty()) {
                continue;
            }

            // 开始爆破（使用线程池）
            callback.onStatusUpdate("开始弱口令爆破: " + formUrl);
            executorService.submit(() ->
                    bruteForce(formUrl, method, usernameFields, passwordFields, formHtml, callback)
            );
        }
    }

    private boolean isUsernameField(String name, String type) {
        name = name.toLowerCase();
        return USERNAME_FIELDS.stream().anyMatch(name::contains) &&
                !"hidden".equalsIgnoreCase(type);
    }

    private boolean isPasswordField(String name, String type) {
        name = name.toLowerCase();
        return PASSWORD_FIELDS.stream().anyMatch(name::contains) &&
                ("password".equalsIgnoreCase(type) || type == null);
    }

    private void bruteForce(String url, String method,
                            Map<String, String> usernameFields,
                            Map<String, String> passwordFields,
                            String formHtml,
                            ScanResultCallback callback) {

        // 1. 先尝试常见弱口令组合
        boolean commonFound = tryCommonCredentials(url, method, usernameFields, passwordFields, formHtml, callback);
        if (commonFound) {
            callback.onStatusUpdate("常见弱口令测试成功: " + url);
            return;
        }

        // 2. 尝试所有字典组合（带限流）
        for (String username : dictionary) {
            if (!isRunning.get()) return;

            for (String password : dictionary) {
                if (!isRunning.get()) return;

                // 限流控制
                try {
                    requestSemaphore.acquire();
                    activeRequests.incrementAndGet();

                    // 创建参数映射
                    Map<String, String> testParams = new HashMap<>();

                    // 设置用户名字段
                    for (Map.Entry<String, String> entry : usernameFields.entrySet()) {
                        testParams.put(entry.getKey(), username);
                    }

                    // 设置密码字段
                    for (Map.Entry<String, String> entry : passwordFields.entrySet()) {
                        testParams.put(entry.getKey(), password);
                    }

                    // 添加其他字段（保留原始值）
                    List<String> otherInputs = new Html(formHtml).xpath("//input").all();
                    StringBuilder paramDebug = new StringBuilder("提交参数: ");
                    String submitName = null;
                    String submitValue = null;

                    for (String inputHtml : otherInputs) {
                        Html input = new Html(inputHtml);
                        String name = input.xpath("//input/@name").get();
                        String type = input.xpath("//input/@type").get();
                        String value = input.xpath("//input/@value").get();

                        if (name != null && !name.isEmpty() &&
                                !testParams.containsKey(name) &&
                                !"submit".equalsIgnoreCase(type)) {
                            String paramValue = value != null ? value : "";
                            testParams.put(name, paramValue);
                            paramDebug.append(name).append("=").append(paramValue).append("&");
                        }

                        // 识别提交按钮
                        if ("submit".equalsIgnoreCase(type) && name != null) {
                            submitName = name;
                            submitValue = value != null ? value : "Submit";
                        }
                    }

                    // 添加提交按钮
                    if (submitName != null) {
                        testParams.put(submitName, submitValue);
                        paramDebug.append(submitName).append("=").append(submitValue).append("&");
                    } else {
                        submitName = "submit";
                        submitValue = "Login";
                        testParams.put(submitName, submitValue);
                        paramDebug.append(submitName).append("=").append(submitValue).append("&");
                        callback.onStatusUpdate("未找到提交按钮，使用默认值");
                    }

                    // 打印完整的请求参数
                    if (paramDebug.length() > 0) {
                        paramDebug.deleteCharAt(paramDebug.length() - 1); // 移除末尾的&
                    }
                    paramDebug.append(" | 用户名字段: ").append(String.join(", ", usernameFields.keySet()))
                            .append(" | 密码字段: ").append(String.join(", ", passwordFields.keySet()))
                            .append(" | 用户名: ").append(username).append(" | 密码: ").append(password);

                    callback.onStatusUpdate(paramDebug.toString());
                    callback.onStatusUpdate("尝试登录: " + username + ":" + password);

                    // 发送请求
                    String response = sendHttpRequest(url, method, testParams);

                    // 检查是否登录成功
                    if (isLoginSuccessful(response)) {
                        String details = "发现弱口令漏洞!\nURL: " + url +
                                "\n用户名: " + username +
                                "\n密码: " + password;
                        callback.onVulnerabilityFound("弱口令", details);
                        return; // 找到一个即停止
                    }

                    // 添加延迟避免请求过快
                    Thread.sleep(200);

                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    return;
                } catch (Exception e) {
                    callback.onError("请求失败: " + e.getMessage());
                } finally {
                    requestSemaphore.release();
                    activeRequests.decrementAndGet();
                }
            }
        }

        callback.onStatusUpdate("弱口令爆破完成，未发现漏洞: " + url);
    }

    private boolean tryCommonCredentials(String url, String method,
                                         Map<String, String> usernameFields,
                                         Map<String, String> passwordFields,
                                         String formHtml,
                                         ScanResultCallback callback) {
        for (String credential : COMMON_CREDENTIALS) {
            if (!isRunning.get()) return false;

            String[] parts = credential.split(":");
            if (parts.length != 2) continue;

            String username = parts[0];
            String password = parts[1];

            try {
                // 创建参数映射
                Map<String, String> testParams = new HashMap<>();

                // 设置用户名字段
                for (Map.Entry<String, String> entry : usernameFields.entrySet()) {
                    testParams.put(entry.getKey(), username);
                }

                // 设置密码字段
                for (Map.Entry<String, String> entry : passwordFields.entrySet()) {
                    testParams.put(entry.getKey(), password);
                }

                // 添加其他字段
                List<String> otherInputs = new Html(formHtml).xpath("//input").all();
                StringBuilder paramDebug = new StringBuilder("[常见密码]提交参数: ");
                String submitName = null;
                String submitValue = null;

                for (String inputHtml : otherInputs) {
                    Html input = new Html(inputHtml);
                    String name = input.xpath("//input/@name").get();
                    String type = input.xpath("//input/@type").get();
                    String value = input.xpath("//input/@value").get();

                    if (name != null && !name.isEmpty() &&
                            !testParams.containsKey(name) &&
                            !"submit".equalsIgnoreCase(type)) {
                        String paramValue = value != null ? value : "";
                        testParams.put(name, paramValue);
                        paramDebug.append(name).append("=").append(paramValue).append("&");
                    }

                    // 识别提交按钮
                    if ("submit".equalsIgnoreCase(type) && name != null) {
                        submitName = name;
                        submitValue = value != null ? value : "Submit";
                    }
                }

                // 添加提交按钮
                if (submitName != null) {
                    testParams.put(submitName, submitValue);
                    paramDebug.append(submitName).append("=").append(submitValue).append("&");
                } else {
                    submitName = "submit";
                    submitValue = "Login";
                    testParams.put(submitName, submitValue);
                    paramDebug.append(submitName).append("=").append(submitValue).append("&");
                    callback.onStatusUpdate("[常见密码]未找到提交按钮，使用默认值");
                }

                // 打印完整的请求参数
                if (paramDebug.length() > 0) {
                    paramDebug.deleteCharAt(paramDebug.length() - 1); // 移除末尾的&
                }
                paramDebug.append(" | 用户名字段: ").append(String.join(", ", usernameFields.keySet()))
                        .append(" | 密码字段: ").append(String.join(", ", passwordFields.keySet()))
                        .append(" | 用户名: ").append(username).append(" | 密码: ").append(password);

                callback.onStatusUpdate(paramDebug.toString());
                callback.onStatusUpdate("[常见密码]尝试登录: " + username + ":" + password);

                // 发送请求
                String response = sendHttpRequest(url, method, testParams);

                // 检查是否登录成功
                if (isLoginSuccessful(response)) {
                    String details = "发现弱口令漏洞!\nURL: " + url +
                            "\n用户名: " + username +
                            "\n密码: " + password;
                    callback.onVulnerabilityFound("弱口令", details);
                    return true;
                }

                // 添加延迟
                Thread.sleep(200);

            } catch (Exception e) {
                callback.onError("常见密码测试失败: " + e.getMessage());
            }
        }
        return false;
    }

    private boolean isLoginSuccessful(String response) {
        if(response.contains("退出登录") ||
                response.contains("Logout") ||
                response.contains("Welcome") ||
                response.contains("个人中心") ||
                response.contains("控制面板") ||
                response.contains("Dashboard") ||
                response.contains("login success")||
                response.contains("My Account")
        ){
            return true;
        }
        return false;
    }
}